import os
import numpy as np
import pandas as pd
import nibabel as nib
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from skimage.transform import resize
from tqdm import tqdm

# ================== 配置参数 ==================
config = {
    "data_root": "./datasets/single_model",
    "batch_size": 8,
    "num_epochs": 30,
    "learning_rate": 1e-4,
    "input_size": (64, 64, 64),
    "num_classes": 4,
    "seed": 42
}

# ================== 设置随机种子 ==================
torch.manual_seed(config["seed"])
np.random.seed(config["seed"])

# ================== 数据加载器 ==================
class MedicalDataset(Dataset):
    def __init__(self, mode="train", transform=None):
        super().__init__()
        self.mode = mode
        self.transform = transform
        self.img_dir = os.path.join(config["data_root"], mode, "images")
        self.mask_dir = os.path.join(config["data_root"], mode, "masks")
        self.label_df = pd.read_csv(os.path.join(config["data_root"], mode, "label.csv"))
        
    def __len__(self):
        return len(self.label_df)

    def load_nii(self, path):
        return nib.load(path).get_fdata().astype(np.float32)
    
    def __getitem__(self, idx):
        p_id = self.label_df.iloc[idx]['p_id']
        label = self.label_df.iloc[idx]['label']-1
        
        # 加载图像和掩膜
        img_path = os.path.join(self.img_dir, f"{p_id}_T2_axi_000.nii.gz")
        mask_path = os.path.join(self.mask_dir, f"{p_id}_T2_axi_roi.nii.gz")
        img = self.load_nii(img_path)
        mask = self.load_nii(mask_path)
        
        # ROI裁剪
        coords = np.where(mask > 0)
        if len(coords[0]) == 0:
            roi = np.zeros(config["input_size"], dtype=np.float32)
        else:
            z_min, z_max = np.min(coords[0]), np.max(coords[0])
            y_min, y_max = np.min(coords[1]), np.max(coords[1])
            x_min, x_max = np.min(coords[2]), np.max(coords[2])
            roi = img[z_min:z_max+1, y_min:y_max+1, x_min:x_max+1]
        
        # 调整尺寸和标准化
        roi = resize(roi, config["input_size"], order=1, preserve_range=True)
        roi = (roi - roi.mean()) / (roi.std() + 1e-8)
        
        # 转换为Tensor
        roi_tensor = torch.FloatTensor(roi).unsqueeze(0)  # [1, D, H, W]
        label_tensor = torch.LongTensor([label])
        
        return roi_tensor, label_tensor

def get_dataloader(mode="train"):
    dataset = MedicalDataset(mode=mode)
    return DataLoader(
        dataset,
        batch_size=config["batch_size"],
        shuffle=(mode == "train"),
        num_workers=4,
        pin_memory=True
    )

# ================== 3D CNN模型 ==================
class CNN3D(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv3d(1, 32, kernel_size=3, padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU(),
            nn.MaxPool3d(2),  # 32x32x32
            
            nn.Conv3d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm3d(64),
            nn.ReLU(),
            nn.MaxPool3d(2),  # 16x16x16
            
            nn.Conv3d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm3d(128),
            nn.ReLU(),
            nn.MaxPool3d(2),  # 8x8x8
        )
        
        self.classifier = nn.Sequential(
            nn.Linear(128 * 8**3, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, config["num_classes"])
        )
        
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

# ================== 训练流程 ==================
def train_model():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    
    # 初始化模型
    model = CNN3D().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=config["learning_rate"])
    
    # 数据加载器
    train_loader = get_dataloader("train")
    val_loader = get_dataloader("val")
    
    best_acc = 0.0
    for epoch in range(config["num_epochs"]):
        # 训练阶段
        model.train()
        running_loss = 0.0
        with tqdm(train_loader, unit="batch") as tepoch:
            for inputs, labels in tepoch:
                tepoch.set_description(f"Epoch {epoch+1}")
                
                inputs = inputs.to(device)
                labels = labels.to(device).view(-1)
                
                optimizer.zero_grad()
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()
                
                running_loss += loss.item()
                tepoch.set_postfix(loss=loss.item())
        
        # 验证阶段
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs = inputs.to(device)
                labels = labels.to(device).view(-1)
                
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        
        val_acc = 100 * correct / total
        print(f"Val Acc: {val_acc:.2f}%")
        
        # 保存最佳模型
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), "best_model_3d_cnn.pth")
            print(f"New best model saved with acc {best_acc:.2f}%")
    
    print("Training complete")
    return model

import matplotlib.pyplot as plt

def visualize_sample(dataset, sample_idx=0):
    """可视化3D医学图像样本的三个正交平面视图"""
    # 获取样本数据
    image_tensor, label_tensor = dataset[sample_idx]
    image = image_tensor.squeeze().numpy()  # 移除通道维度 (C, D, H, W) -> (D, H, W)
    label = label_tensor.item() + 1         # 恢复原始标签 (0-based -> 1-based)

    # 计算中心切片索引
    depth, height, width = image.shape
    axial_slice = image[depth//2, :, :]      # 轴向切片（水平面）
    sagittal_slice = image[:, height//2, :]  # 矢状切片（侧面）
    coronal_slice = image[:, :, width//2]    # 冠状切片（正面）

    # 创建可视化图像
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    slices = [axial_slice, sagittal_slice, coronal_slice]
    titles = ['Axial View', 'Sagittal View', 'Coronal View']

    for ax, img, title in zip(axes, slices, titles):
        # 显示图像并调整方向
        ax.imshow(img.T, cmap='gray', origin='lower')
        ax.set_title(title)
        ax.axis('off')

    plt.suptitle(f'Patient Sample (Class {label})', y=0.95)
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    # train_model()
    # 创建数据集实例
    train_dataset = MedicalDataset(mode="train")
    
    # 可视化第一个样本
    print("Visualizing sample:")
    visualize_sample(train_dataset, sample_idx=0)
